import torch.nn as nn

class Heatmap_Classify_Loss(nn.Module):
    def __init__(self):
        super(Heatmap_Classify_Loss, self).__init__()
        self.heatmap_loss_fun = nn.MSELoss()
        self.classify_loss_fun = nn.CrossEntropyLoss()
    def forward(self,input,label):
        # print(input[1].shape,label[1].shape)
        heatmap_loss = self.heatmap_loss_fun(input[0],label[0])
        classify_loss_1 = self.classify_loss_fun(input[1],label[1].long())
        classify_loss_2 = self.classify_loss_fun(input[2],label[2].long())
        # return heatmap_loss
        return heatmap_loss+classify_loss_1+classify_loss_2
